#include "parser.h"
#include "Diagnosis.h"
#include <cassert>
#include <iostream>

std::shared_ptr<ProgramAstNode> Parser::Parse() {
    auto node = std::make_shared<ProgramAstNode>(Lex.CurrentToken);
    while (Lex.CurrentToken->Kind != TokenKind::Eof) {
        node->Funcs.push_back(ParseFunction());
    }
    return node;
}

std::shared_ptr<AstNode> Parser::ParseFunction() {
    auto node = std::make_shared<FunctionNode>(Lex.CurrentToken);
    Locals = &(node->LocalVars);
    LocalsMap.clear();

    auto ty = ParseDeclarationSpec();       // 解析函数返回的基本类型
    std::shared_ptr<Token> funcNameToken;
    ty = ParseDeclarator(ty, funcNameToken);    // 为返回基本类型加指针，同时获取函数名token

    // 函数名
    node->FuncName = funcNameToken->Content;
    node->Ty = ty;  
    
    auto funcType = std::dynamic_pointer_cast<FunctionType>(ty);
    // 保持node->Params和node->LocalVars中参数顺序一样
    for (auto it = funcType->Params.rbegin(); it != funcType->Params.rend(); ++it) {
        node->Params.push_front(MakeLocalVar((*it)->tk->Content, (*it)->Ty));
    }
    
    // 函数体
    Lex.ExceptToken(TokenKind::LBrace);
    while (Lex.CurrentToken->Kind != TokenKind::RBrace) {
        node->Statms.push_back(ParseStatement());
    }
    Lex.ExceptToken(TokenKind::RBrace);
    return node;
}

std::shared_ptr<Type> Parser::ParseDeclarationSpec() {
    if (Lex.CurrentToken->Kind == TokenKind::Char) {
        Lex.GetNextToken();
        return Type::CharType;
    } else if (Lex.CurrentToken->Kind == TokenKind::Short) {
        Lex.GetNextToken();
        return Type::ShortType;
    } else if (Lex.CurrentToken->Kind == TokenKind::Int) {
        Lex.GetNextToken();
        return Type::IntType;
    } else if (Lex.CurrentToken->Kind == TokenKind::Long) {
        Lex.GetNextToken();
        return Type::LongType;
    }
    std::string content(Lex.CurrentToken->Content);
    DiagErr(Lex.SourceCode, Lex.CurrentToken->Location.Col,
            "Current Type '%s' is not supported!(line: %d)\n", 
            content.data(), Lex.CurrentToken->Location.Line); 
}

std::shared_ptr<Type> Parser::ParseDeclarator(std::shared_ptr<Type> baseType, std::shared_ptr<Token> &token) {
    auto ty = baseType;
    while (Lex.CurrentToken->Kind == TokenKind::Star) {
        ty = std::make_shared<PointerType>(ty);
        Lex.GetNextToken();
    }
    if (Lex.CurrentToken->Kind != TokenKind::Identifier) {
        std::string content(Lex.CurrentToken->Content);
        DiagErr(Lex.SourceCode, Lex.CurrentToken->Location.Col,
            "a Variable Name is Excepted!, but got '%s' (line: %d)\n", 
            content.data(), Lex.Line); 
    }

    token = Lex.CurrentToken;
    Lex.ExceptToken(TokenKind::Identifier);
    return ParseTypeSuffix(ty);
}

std::shared_ptr<Type> Parser::ParseTypeSuffix(std::shared_ptr<Type> baseType) {
    if (Lex.CurrentToken->Kind == TokenKind::LParent) {
        Lex.GetNextToken();
        auto funcType = std::make_shared<FunctionType>(baseType);
        if (Lex.CurrentToken->Kind != TokenKind::RParent) {
            std::shared_ptr<Token> ParamNametoken;
            auto ty = ParseDeclarator(ParseDeclarationSpec(), ParamNametoken);
            auto pa = std::make_shared<Param>();

            pa->tk = ParamNametoken;
            pa->Ty = ty;
            funcType->Params.push_back(pa);

            while (Lex.CurrentToken->Kind != TokenKind::RParent) {
                Lex.ExceptToken(TokenKind::Comma); 
                auto ty = ParseDeclarator(ParseDeclarationSpec(), ParamNametoken);
                auto pa = std::make_shared<Param>();

                pa->tk = ParamNametoken;
                pa->Ty = ty;
                funcType->Params.push_back(pa);
            }
        }
        Lex.ExceptToken(TokenKind::RParent);
        return funcType;
    } 
    else if (Lex.CurrentToken->Kind == TokenKind::LBacket) {
        Lex.ExceptToken(TokenKind::LBacket);
        int num = Lex.CurrentToken->Value;
        Lex.ExceptToken(TokenKind::Num);
        Lex.ExceptToken(TokenKind::RBacket);
        auto ty = ParseTypeSuffix(baseType);
        return std::make_shared<ArrayType>(ty, num);
    }
    else {
        return baseType;
    }
}

std::shared_ptr<AstNode> Parser::ParseStatement() {
    if (Lex.CurrentToken->Kind == TokenKind::If) {
        auto node = std::make_shared<IfStatmNode>(Lex.CurrentToken);
        Lex.GetNextToken();
        Lex.ExceptToken(TokenKind::LParent);
        node->Cond = ParseExpr();
        Lex.ExceptToken(TokenKind::RParent);
        node->Then = ParseStatement();
        if (Lex.CurrentToken->Kind == TokenKind::Else) {
            Lex.GetNextToken();
            node->Else = ParseStatement();
        }
        return node;
    } 
    else if (Lex.CurrentToken->Kind == TokenKind::While) {
        auto node = std::make_shared<WhileStatmNode>(Lex.CurrentToken);
        Lex.GetNextToken();
        Lex.ExceptToken(TokenKind::LParent);
        node->Cond = ParseExpr();
        Lex.ExceptToken(TokenKind::RParent);
        node->Then = ParseStatement();
        return node;
    }
    else if (Lex.CurrentToken->Kind == TokenKind::For) {
        auto node = std::make_shared<ForStatmNode>(Lex.CurrentToken);
        Lex.GetNextToken();
        Lex.ExceptToken(TokenKind::LParent);
        if (Lex.CurrentToken->Kind != TokenKind::Semicolon) {   // 空语句
            node->Init = ParseExpr();
        }     
        Lex.ExceptToken(TokenKind::Semicolon);
        if (Lex.CurrentToken->Kind != TokenKind::Semicolon) {   // 空语句
            node->Cond = ParseExpr();
        }     
        Lex.ExceptToken(TokenKind::Semicolon);
            if (Lex.CurrentToken->Kind != TokenKind::Semicolon) {   // 空语句
            node->Inc = ParseExpr();
        }     
        Lex.ExceptToken(TokenKind::RParent);
        node->Then = ParseStatement();
        return node;
    }
    else if (Lex.CurrentToken->Kind == TokenKind::Return) {
        
        auto node = std::make_shared<ReturnStatmNode>(Lex.CurrentToken);
        Lex.GetNextToken();
        node->Lhs = ParseExpr();
        Lex.ExceptToken(TokenKind::Semicolon);
        return node;
    }
    else if (Lex.CurrentToken->Kind == TokenKind::LBrace) {
        auto node = std::make_shared<BlockStatmNode>(Lex.CurrentToken);
        Lex.GetNextToken();
        while (Lex.CurrentToken->Kind != TokenKind::RBrace) {
            node->Statms.push_back(ParseStatement());
        }
        Lex.ExceptToken(TokenKind::RBrace);
        return node;   
    }
    else if (IsTypeName()) {
        /// 支持 int a; int a = 3; int a, b = 3;
        auto deaclraStatmNode = std::make_shared<DeclarationStatmNode>(Lex.CurrentToken);
        auto ty = ParseDeclarationSpec();
        int varCount = 0;                                               // 声明语句 声明的变量数量
        while (Lex.CurrentToken->Kind != TokenKind::Semicolon) {
            if (varCount > 0) {
                Lex.ExceptToken(TokenKind::Comma);              // 函数参数中间的','
            }
            std::shared_ptr<Token> tok;
            ty = ParseDeclarator(ty, tok);

            auto var = MakeLocalVar(tok->Content, ty);
            varCount++;
            if (Lex.CurrentToken->Kind != TokenKind::Assign) {
                continue;
            }
            auto assignNode = std::make_shared<AssignExprNode>(Lex.CurrentToken);

            Lex.GetNextToken();

            assignNode->Lhs = MakeVarNode(var, tok);
            assignNode->Rhs = ParseExpr();
            deaclraStatmNode->AssignNodes.push_back(assignNode);
        }
        Lex.ExceptToken(TokenKind::Semicolon);
        return deaclraStatmNode;
    }
    else {
        auto node = std::make_shared<ExprStatmNode>(Lex.CurrentToken);
        if (Lex.CurrentToken->Kind != TokenKind::Semicolon) {   // 单分号 空语句
            node->Lhs = ParseExpr();
        }
        Lex.ExceptToken(TokenKind::Semicolon);
        return node;
    }
}

std::shared_ptr<AstNode> Parser::ParseFuncCall(std::shared_ptr<Token> nameTok) {
    auto node = std::make_shared<FuncCallNode>(nameTok);
    node->FuncName = nameTok->Content;
    Lex.ExceptToken(TokenKind::LParent);
    if (Lex.CurrentToken->Kind != TokenKind::RParent) {
        node->Args.push_back(ParseAssignExpr());
        while (Lex.CurrentToken->Kind == TokenKind::Comma) {
            Lex.GetNextToken();
            node->Args.push_back(ParseAssignExpr());
        }
    }
    Lex.ExceptToken(TokenKind::RParent);
    return node;
}

std::shared_ptr<AstNode> Parser::ParseExpr() {
    return ParseAssignExpr(); 
}   

std::shared_ptr<AstNode> Parser::ParseAssignExpr() {
    std::shared_ptr<AstNode> left = ParseEqualExpr();
    if (Lex.CurrentToken->Kind == TokenKind::Assign) {
        auto node = std::make_shared<AssignExprNode>(Lex.CurrentToken);
        Lex.GetNextToken();
        node->Lhs = left;
        node->Rhs = ParseAssignExpr();
        return node;
    }
    return left;
}

std::shared_ptr<AstNode> Parser::ParseEqualExpr() {
    std::shared_ptr<AstNode> left = ParseRelationExpr();
    while (Lex.CurrentToken->Kind == TokenKind::Equal || Lex.CurrentToken->Kind == TokenKind::PipeEqual) {
        BinaryOperator op = BinaryOperator::Equal;
        if (Lex.CurrentToken->Kind == TokenKind::PipeEqual) op = BinaryOperator::PipeEqual;
        auto node = std::make_shared<BinaryAstNode>(Lex.CurrentToken);
        
        Lex.GetNextToken();
        node->Lhs = left;
        node->BinOp = op;
        node->Rhs = ParseRelationExpr();
        left = node;
    }
    return left;
}

std::shared_ptr<AstNode> Parser::ParseRelationExpr() {
    std::shared_ptr<AstNode> left = ParseAddExpr();
    while (Lex.CurrentToken->Kind == TokenKind::Greater 
        || Lex.CurrentToken->Kind == TokenKind::GreaterEqual
        || Lex.CurrentToken->Kind == TokenKind::Lesser
        || Lex.CurrentToken->Kind == TokenKind::LesserEqual) {
        
        BinaryOperator op = BinaryOperator::Greater;
        if (Lex.CurrentToken->Kind == TokenKind::GreaterEqual) op = BinaryOperator::GreaterEqual;
        else if (Lex.CurrentToken->Kind == TokenKind::Lesser) op = BinaryOperator::Lesser;
        else if (Lex.CurrentToken->Kind == TokenKind::LesserEqual) op = BinaryOperator::LesserEqual;
        
        auto node = std::make_shared<BinaryAstNode>(Lex.CurrentToken);
        
        Lex.GetNextToken();
        node->Lhs = left;
        node->BinOp = op;
        node->Rhs = ParseAddExpr();
        left = node;
    }
    return left;
}

std::shared_ptr<AstNode> Parser::ParseAddExpr() {
    std::shared_ptr<AstNode> left = ParseMultiExpr();
    while (Lex.CurrentToken->Kind == TokenKind::Plus || Lex.CurrentToken->Kind == TokenKind::Minus) {
        
        auto tok = Lex.CurrentToken;
        Lex.GetNextToken();

        auto right = ParseMultiExpr();
        left->Accept(TypeVisitor::Visitor());
        right->Accept(TypeVisitor::Visitor());

        BinaryOperator op = BinaryOperator::Add;
        if (tok->Kind == TokenKind::Plus) {
            if ((left->Ty->IsPointerType() || left->Ty->IsArrayType()) && right->Ty->IsIntegerType()) {
                op = BinaryOperator::PtrAdd;
            } else if (left->Ty->IsIntegerType() && (right->Ty->IsPointerType() || right->Ty->IsArrayType())) {
                swap(left, right);
                op = BinaryOperator::PtrAdd;
            } else if (left->Ty->IsIntegerType() && right->Ty->IsIntegerType()) {
                op = BinaryOperator::Add;
            } else {
                DiagErr(Lex.SourceCode, Lex.CurrentToken->Location.Col, "invalid add operation(%d)\n", Lex.CurrentToken->Location.Line);
            }
        } else {
            if ((left->Ty->IsPointerType() || left->Ty->IsArrayType()) && right->Ty->IsIntegerType()) {
                op = BinaryOperator::PtrSub;
            } else if (left->Ty->IsIntegerType() && right->Ty->IsIntegerType()) {
                op = BinaryOperator::Sub;
            } else if ((left->Ty->IsPointerType() || left->Ty->IsArrayType())
                    && (right->Ty->IsPointerType() || right->Ty->IsArrayType())) {
                op = BinaryOperator::PtrDiff;
            } else {
                DiagErr(Lex.SourceCode, Lex.CurrentToken->Location.Col, "invalid sub operation(%d)\n", Lex.CurrentToken->Location.Line);
            }        
        }
        
        auto node = std::make_shared<BinaryAstNode>(tok); 
        node->BinOp = op;
        node->Lhs = left;
        node->Rhs = right;
        left = node;
    }
    return left;
}

std::shared_ptr<AstNode> Parser::ParseMultiExpr() {
    std::shared_ptr<AstNode> left = ParseUnaryExpr();
    while (Lex.CurrentToken->Kind == TokenKind::Star || Lex.CurrentToken->Kind == TokenKind::Slash) {
        BinaryOperator op = BinaryOperator::Mul;
        if (Lex.CurrentToken->Kind == TokenKind::Slash) op = BinaryOperator::Div;
        
        auto node = std::make_shared<BinaryAstNode>(Lex.CurrentToken);

        Lex.GetNextToken();
        node->BinOp = op;
        node->Lhs = left;
        node->Rhs = ParseUnaryExpr();
        left = node;
    }
    return left;
}

std::shared_ptr<AstNode> Parser::ParseUnaryExpr()
{
    if (Lex.CurrentToken->Kind == TokenKind::Plus || Lex.CurrentToken->Kind == TokenKind::Minus
    || Lex.CurrentToken->Kind == TokenKind::Star || Lex.CurrentToken->Kind == TokenKind::Amp) {
        switch (Lex.CurrentToken->Kind) {
            case TokenKind::Plus: {
                auto node = std::make_shared<UnaryAstNode>(Lex.CurrentToken);
                node->UnaryOp = UnaryOperator::Plus;
                Lex.GetNextToken();
                node->Lhs = ParseUnaryExpr();
                return node;
            }
            case TokenKind::Minus: {
                auto node = std::make_shared<UnaryAstNode>(Lex.CurrentToken);
                node->UnaryOp = UnaryOperator::Minus;
                Lex.GetNextToken();
                node->Lhs = ParseUnaryExpr();
                return node;
            }
            case TokenKind::Star: {
                auto node = std::make_shared<UnaryAstNode>(Lex.CurrentToken);
                node->UnaryOp = UnaryOperator::Deref;
                Lex.GetNextToken();
                node->Lhs = ParseUnaryExpr();
                return node;
            }
            case TokenKind::Amp: {
                auto node = std::make_shared<UnaryAstNode>(Lex.CurrentToken);
                node->UnaryOp = UnaryOperator::Amp;
                Lex.GetNextToken();
                node->Lhs = ParseUnaryExpr();
                return node;
            }
            default:
            assert(0);  
                break;
        } // switch
    } // if
    
    return ParsePostSuffixExpr();
}

std::shared_ptr<AstNode> Parser::ParsePostSuffixExpr() {
    auto left = ParsePrimaryExpr();
    while (true) {
        // if (Lex.CurrentToken->Kind == TokenKind::LParent) {
        //     return ParseFuncCall(left->Tok);
        // } else 
        if (Lex.CurrentToken->Kind == TokenKind::LBacket) {
            /// a[3] => *(a+3) or a[3][4] => *(*(a+3)+4)
            auto ptrAddNode = std::make_shared<BinaryAstNode>(left->Tok);     // 先创建PtrAdd二元节点
            Lex.ExceptToken(TokenKind::LBacket);
            
            ptrAddNode->Lhs = left;
            ptrAddNode->BinOp = BinaryOperator::PtrAdd;
            ptrAddNode->Rhs = ParseExpr();
            Lex.ExceptToken(TokenKind::RBacket);

            auto starNode = std::make_shared<UnaryAstNode>(left->Tok);          // 再创建*一元节点
            starNode->UnaryOp = UnaryOperator::Deref;
            starNode->Lhs = ptrAddNode;
            left = starNode;                                                    // 可能遇到a[3][4]
        } else {
            break;
        }
    }
    return left;
}

std::shared_ptr<AstNode> Parser::ParsePrimaryExpr() {
    if (Lex.CurrentToken->Kind == TokenKind::LParent) {
        Lex.GetNextToken();
        auto node = ParseExpr();
        Lex.ExceptToken(TokenKind::RParent);     // 跳过右括号
        return node;
    }
    else if (Lex.CurrentToken->Kind == TokenKind::Identifier) {
        
        Lex.SaveStatus();                                       // 保存当前状态
        Lex.GetNextToken();                                     // 向前查看一个状态
        if (Lex.CurrentToken->Kind == TokenKind::LParent) {     // 查看是否为左括号，是的话就进行函数调用
            Lex.LoadStatus();                                   // 加载保存的状态
            auto tok = Lex.CurrentToken;
            Lex.ExceptToken(TokenKind::Identifier);
            return ParseFuncCall(tok);
        }
        Lex.LoadStatus();                                       // 加载保存的状态
        
        std::shared_ptr<Var> obj = FindLocalVar(Lex.CurrentToken->Content);
        if (!obj) {
            DiagErr(Lex.SourceCode, Lex.CurrentToken->Location.Col, 
            "undefined variable(line: %d)", Lex.CurrentToken->Location.Line);
        }
        auto node = MakeVarNode(obj, Lex.CurrentToken);
        Lex.GetNextToken();
        return node;
    } else if (Lex.CurrentToken->Kind == TokenKind::Num) {
        auto node = std::make_shared<ConstantAstNode>(Lex.CurrentToken);
        node->Value = Lex.CurrentToken->Value;
        Lex.GetNextToken();
        return node; 
    } else if (Lex.CurrentToken->Kind == TokenKind::Sizeof) {
        auto node = std::make_shared<SizeofNode>(Lex.CurrentToken);
        Lex.GetNextToken();
        node->Lhs = ParseUnaryExpr();
        // node->Ty 
        return node;
    } else {
        std::string content(Lex.CurrentToken->Content);
        DiagErr(Lex.SourceCode, Lex.CurrentToken->Location.Col, 
                "node '%s' is not supported (line: %d)", 
                content.data(),Lex.CurrentToken->Location.Line);
    }
}

bool Parser::IsTypeName() {
    if (Lex.CurrentToken->Kind == TokenKind::Char
    || Lex.CurrentToken->Kind == TokenKind::Short
    || Lex.CurrentToken->Kind == TokenKind::Int
    || Lex.CurrentToken->Kind == TokenKind::Long){
        return true;
    }
    return false;
}

std::shared_ptr<AstNode> Parser::MakeVarNode(std::shared_ptr<Var> var, std::shared_ptr<Token> tok)
{
    auto node = std::make_shared<VarAstNode>(tok);
    node->VarObj = var;
    node->Ty = var->Ty;
    return node;
}

std::shared_ptr<Var> Parser::FindLocalVar(std::string_view name) {
    if (LocalsMap.find(name) != LocalsMap.end()) {
        return LocalsMap[name];
    }
    return nullptr;
}

std::shared_ptr<Var> Parser::MakeLocalVar(std::string_view name, std::shared_ptr<Type> ty) {
    auto obj = std::make_shared<Var>();
    obj->Ty = ty;
    obj->Name = name;
    obj->Offset = 0;
    Locals->push_front(obj);
    LocalsMap[name] = obj;
    return obj;
}